#include "socket.h"
#include <algorithm>
#include <string.h> // FD_SET needs memset on some platforms, so we can't use <cstring>
#include <sys/epoll.h>
#include <poll.h>
#include <errno.h>
#include <netinet/tcp.h>
#include <sys/time.h>

#include "timestat.h"

Socket::Socket(__socket_type sockettype):
    m_sockfd(INVALID_SOCKET),
    m_socketType(sockettype)
{
}

Socket::Socket(int sockfd,__socket_type sockettype):
    m_sockfd(sockfd),
    m_socketType(sockettype)
{
}

Socket::Socket(const Socket& other)
{
    if (&other != this)
    {
        m_sockfd = other.m_sockfd;
        m_socketType = other.m_socketType;
    }
}

Socket& Socket::operator = (const Socket& other)
{
    if (&other != this)
    {
        m_sockfd = other.m_sockfd;
        m_socketType = other.m_socketType;
    }

    return *this;
}

Socket::~Socket()
{
}

void Socket::close()
{
    if (m_sockfd != INVALID_SOCKET)
    {
        ::close(m_sockfd);
        m_sockfd = INVALID_SOCKET;
    }
}   
    
void Socket::init(int af)
{
    initSocket(af, m_socketType);
}

int Socket::select(vector<Socket>& readList, 
        vector<Socket>& writeList,
        vector<Socket>& exceptList,
        unsigned int timeout_ms)
{
    int epollSize = readList.size() + writeList.size() + exceptList.size();
    if (epollSize == 0) return 0;

    struct epoll_event eventsIn[epollSize];
    memset(eventsIn, 0, sizeof(eventsIn));
    struct epoll_event* eventLast = eventsIn;
    for (vector<Socket>::iterator it = readList.begin(); it != readList.end(); ++it)
    {
        int sockfd = it->sockfd();
        if (sockfd != INVALID_SOCKET)
        {
            struct epoll_event* e = eventsIn;
            for (; e != eventLast; ++e)
            {
                if (reinterpret_cast<Socket*>(e->data.ptr)->sockfd() == sockfd)
                    break;
            }
            if (e == eventLast)
            {
                e->data.ptr = &(*it);
                ++eventLast;
            }
            e->events |= EPOLLIN;
        }
    }

    for (vector<Socket>::iterator it = writeList.begin(); it != writeList.end(); ++it)
    {
        int sockfd = it->sockfd();
        if (sockfd != INVALID_SOCKET)
        {
            struct epoll_event* e = eventsIn;
            for (; e != eventLast; ++e)
            {
                if (reinterpret_cast<Socket*>(e->data.ptr)->sockfd() == sockfd)
                    break;
            }
            if (e == eventLast)
            {
                e->data.ptr = &(*it);
                ++eventLast;
            }
            e->events |= EPOLLOUT;
        }
    }

    for (vector<Socket>::iterator it = exceptList.begin(); it != exceptList.end(); ++it)
    {
        int sockfd = it->sockfd();
        if (sockfd != INVALID_SOCKET)
        {
            struct epoll_event* e = eventsIn;
            for (; e != eventLast; ++e)
            {
                if (reinterpret_cast<Socket*>(e->data.ptr)->sockfd() == sockfd)
                    break;
            }
            if (e == eventLast)
            {
                e->data.ptr = &(*it);
                ++eventLast;
            }
            e->events |= EPOLLERR;
        }
    }

    epollSize = eventLast - eventsIn;
    int epollfd = epoll_create(epollSize);
    if (epollfd < 0)
    {
        return 0;
    }

    for (struct epoll_event* e = eventsIn; e != eventLast; ++e)
    {
        int sockfd = reinterpret_cast<Socket*>(e->data.ptr)->sockfd();
        if (sockfd != INVALID_SOCKET)
        {
            if (epoll_ctl(epollfd, EPOLL_CTL_ADD, sockfd, e) < 0)
            {
                return 0;
            }
        }
    }

    struct epoll_event eventsOut[epollSize];
    memset(eventsOut, 0, sizeof(eventsOut));

    int rc;
    int remain = timeout_ms;
    do
    {
        CTimeStat tmstat;
        tmstat.AddStep("begin");
        rc = epoll_wait(epollfd, eventsOut, epollSize, remain);
        if (rc < 0)
        {
            tmstat.AddStep("end");
            int waited = tmstat.GetInterval("end","begin",UNIT_USECOND)/1000;
            if (waited < remain)
                remain -= waited;
            else
                break;
        }
    }while (rc < 0 && remain > 0);

    ::close(epollfd);
    if (rc < 0)
        return 0;

    vector<Socket> readyReadList;
    vector<Socket> readyWriteList;
    vector<Socket> readyExceptList;
    for (int n = 0; n < rc; ++n)
    {
        if (eventsOut[n].events & EPOLLERR)
            readyExceptList.push_back(*reinterpret_cast<Socket*>(eventsOut[n].data.ptr));
        if (eventsOut[n].events & EPOLLIN)
            readyReadList.push_back(*reinterpret_cast<Socket*>(eventsOut[n].data.ptr));
        if (eventsOut[n].events & EPOLLOUT)
            readyWriteList.push_back(*reinterpret_cast<Socket*>(eventsOut[n].data.ptr));
    }
    std::swap(readList, readyReadList);
    std::swap(writeList, readyWriteList);
    std::swap(exceptList, readyExceptList);
    return readList.size() + writeList.size() + exceptList.size();
}

int Socket::acceptConnection()
{
    if (m_sockfd == INVALID_SOCKET)
        return INVALID_SOCKET;

    struct sockaddr_storage addr;
    unsigned int len = sizeof(addr);
    memset(&addr,0,len);
    struct sockaddr* pSA = (struct sockaddr*)&addr;
    int fd;
    do
    {
        fd = ::accept(m_sockfd, pSA, &len);
    }
    while (fd == INVALID_SOCKET && errno == EINTR);

    return fd;
}

int Socket::initSocket(int af, int type, int proto)
{
    m_sockfd = ::socket(af, type, proto);
    return m_sockfd;
}

int Socket::connect(const SocketAddress& address)
{
    if (m_sockfd == INVALID_SOCKET)
    {
        initSocket(address.af(), SOCK_STREAM);
    }
    
    int rc = 0;
    do
    {
        struct sockaddr_storage tmp = address.addr();
        rc = ::connect(m_sockfd, (sockaddr*)&tmp, address.length());
    }
    while (rc != 0 && errno == EINTR);

    return rc;
}

int Socket::connect(const SocketAddress& address, unsigned int timeout_ms)
{
    if (m_sockfd == INVALID_SOCKET)
    {
        initSocket(address.af(), SOCK_STREAM);
    }
    
    bool blocking = getBlocking(m_sockfd);
    if (blocking)
    {
        setBlocking(m_sockfd,false);
    }
    int rc = 0;
    struct sockaddr_storage tmp = address.addr();
    rc = ::connect(m_sockfd, (sockaddr*)&tmp, address.length());
    if (rc != 0)
    {
        int err = errno;
        if (err != EINPROGRESS && err != EWOULDBLOCK)
            return err;
        if (timeout_ms > 0 && !poll(timeout_ms, SELECT_READ | SELECT_WRITE | SELECT_ERROR))
            return errno;
    }

    if (blocking)
        setBlocking(m_sockfd,true);
    
    return rc;
}

int Socket::connectNB(const SocketAddress& address)
{
    if (m_sockfd == INVALID_SOCKET)
    {
        initSocket(address.af(), SOCK_STREAM);
    }
    setBlocking(m_sockfd,false);
    struct sockaddr_storage tmp = address.addr();
    int rc = ::connect(m_sockfd, (sockaddr*)&tmp, address.length());
    if (rc != 0)
    {
        int err = errno;
        if (err != EINPROGRESS && err != EWOULDBLOCK)
            return err;
    }
    return rc;
}

int Socket::bind(const SocketAddress& address, bool reuseAddress)
{
    if (m_sockfd == INVALID_SOCKET)
    {
        initSocket(address.af(), SOCK_STREAM);
    }
    if (reuseAddress)
    {
        setReuseAddress(true);
        setReusePort(true);
    }

    struct sockaddr_storage tmp = address.addr();
    int rc = ::bind(m_sockfd, (sockaddr*)&tmp, address.length());
    return rc;
}

int Socket::listen(int backlog)
{
    if (m_sockfd == INVALID_SOCKET)
        return INVALID_SOCKET;
    return ::listen(m_sockfd, backlog);
}

void Socket::shutdownRecv()
{
    if (m_sockfd == INVALID_SOCKET)
        return;

    ::shutdown(m_sockfd, 0);
}

void Socket::shutdownSend()
{
    if (m_sockfd == INVALID_SOCKET)
        return;

    ::shutdown(m_sockfd, 1);
}

void Socket::shutdown()
{
    if (m_sockfd == INVALID_SOCKET)
        return;

    ::shutdown(m_sockfd, 2);
}

int Socket::sendBytes(const char* buffer, unsigned int length, int flags)
{
    if (m_sockfd == INVALID_SOCKET)
        return INVALID_SOCKET;

    int rc = 0;
    do
    {
        rc = ::send(m_sockfd, buffer, length, flags);
    }while (rc < 0 && (errno == EINTR || errno == EAGAIN));

    return rc;
}

int Socket::recvBytes(char* buffer, unsigned int length, int flags)
{
    if (m_sockfd == INVALID_SOCKET )
        return INVALID_SOCKET;

    if ( buffer == NULL || length == 0 )
        return 0;

    int rc = 0;
    do
    {
        rc = ::recv(m_sockfd, buffer, length, flags);
    }while (rc < 0 && (errno == EINTR || errno == EAGAIN));

    return rc;
}

int Socket::sendTo(const char* buffer, unsigned int length, const SocketAddress& address, int flags)
{
    if (m_sockfd == INVALID_SOCKET)
        return INVALID_SOCKET;
    int rc = 0;
    struct sockaddr_storage tmp = address.addr();
    do
    {
        rc = ::sendto(m_sockfd, buffer, length, flags, (sockaddr*)&tmp, address.length());
    }while (rc < 0 && (errno == EINTR || errno == EAGAIN ));

    return rc;
}

int Socket::recvFrom(char* buffer, unsigned int length, SocketAddress& address, int flags)
{
    if (m_sockfd == INVALID_SOCKET)
        return INVALID_SOCKET;
    struct sockaddr_storage addr;
    unsigned int saLen = sizeof(addr);
    memset(&addr,0,saLen);
    struct sockaddr* pSA = (struct sockaddr*)&addr;
    int rc;
    do
    {
        rc = ::recvfrom(m_sockfd, buffer, length, flags, pSA, &saLen);
    }while (rc < 0 && (errno == EINTR || errno == EAGAIN ));
    
    if (rc >= 0)
    {
        address = SocketAddress(addr);
    }
    
    return rc;
}

void Socket::sendUrgent(unsigned char data)
{
    if (m_sockfd == INVALID_SOCKET)
        return;
    ::send(m_sockfd, reinterpret_cast<const char*>(&data), sizeof(data), MSG_OOB);
}

int Socket::available(int fd)
{
    int result;
    ::ioctl(fd, FIONREAD, &result);
    return result;
}

bool Socket::poll(unsigned int timeout_ms, int mode)
{
    if (m_sockfd == INVALID_SOCKET)
        return false;
    int epollfd = epoll_create(1);
    if (epollfd < 0)
    {
        return false;
    }

    struct epoll_event evin;
    memset(&evin, 0, sizeof(evin));

    if (mode & SELECT_READ)
        evin.events |= EPOLLIN;
    if (mode & SELECT_WRITE)
        evin.events |= EPOLLOUT;
    if (mode & SELECT_ERROR)
        evin.events |= EPOLLERR;

    if (epoll_ctl(epollfd, EPOLL_CTL_ADD,m_sockfd, &evin) < 0)
    {
        return false;
    }

    unsigned int remain = timeout_ms;
    int rc;
    do
    {
        struct epoll_event evout;
        memset(&evout, 0, sizeof(evout));

        CTimeStat tmstat;
        tmstat.AddStep("begin");
        rc = epoll_wait(epollfd, &evout, 1, remain);
        if (rc < 0 && errno == EINTR)
        {
            tmstat.AddStep("end");
            long waited = tmstat.GetInterval("end","start",UNIT_USECOND)/1000;
            if (waited < remain)
                remain -= waited;
            else
                break;
        }
    }
    while (rc < 0 && errno == EINTR);

    ::close(epollfd);
    
    return rc > 0; 
}

void Socket::setSendBufferSize(unsigned int size)
{
    setOption(SOL_SOCKET, SO_SNDBUF, &size,sizeof(size));
}

unsigned int Socket::getSendBufferSize()
{
    unsigned int result = 0;
    unsigned int len = sizeof(result);
    getOption(SOL_SOCKET, SO_SNDBUF, &result,len);
    return result;
}

void Socket::setRecvBufferSize(unsigned int size)
{
    setOption(SOL_SOCKET, SO_RCVBUF, &size,sizeof(size));
}

unsigned int Socket::getRecvBufferSize()
{
    unsigned int result = 0;
    unsigned int len = sizeof(result);
    getOption(SOL_SOCKET, SO_RCVBUF, &result,len);
    return result;
}

void Socket::setSendTimeout(timeval& val)
{
    setOption(SOL_SOCKET, SO_SNDTIMEO, &val,sizeof(val));
}

void Socket::getSendTimeout(timeval& val)
{
    unsigned int len = sizeof(val);
    getOption(SOL_SOCKET, SO_SNDTIMEO, &val,len);
}

void Socket::setRecvTimeout(timeval& val)
{
    setOption(SOL_SOCKET, SO_RCVTIMEO, &val,sizeof(val));
}

void Socket::getRecvTimeout(timeval& val)
{
    unsigned int len = sizeof(val);
    getOption(SOL_SOCKET, SO_RCVTIMEO, &val,len);
}

SocketAddress Socket::address()
{
    if (m_sockfd == INVALID_SOCKET)
        return SocketAddress();
    
    struct sockaddr_storage addr;
    unsigned int saLen = sizeof(addr);
    memset(&addr,0,saLen);
    struct sockaddr* pSA = (struct sockaddr*)&addr;
    int rc = ::getsockname(m_sockfd, pSA, &saLen);
    if (rc == 0)
        return SocketAddress(addr);

    return SocketAddress();
}

SocketAddress Socket::peerAddress()
{
    return SocketAddress::peerAddress(m_sockfd);
}

void Socket::setOption(int level, int option, const void* value, int length)
{
    if (m_sockfd == INVALID_SOCKET)
        return;
    ::setsockopt(m_sockfd, level, option, reinterpret_cast<const char*>(value), length);
}

void Socket::setOption(int sockfd,int level, int option, const void* value, int length)
{
    if (sockfd == INVALID_SOCKET)
        return;
    ::setsockopt(sockfd, level, option, reinterpret_cast<const char*>(value), length);
}

void Socket::getOption(int level, int option, void* value, unsigned int& length)
{
    if (m_sockfd == INVALID_SOCKET)
        return;
    
    ::getsockopt(m_sockfd, level, option, reinterpret_cast<char*>(value), &length);
}

void Socket::getOption(int sockfd,int level, int option, void* value,unsigned int& length)
{
    if (sockfd == INVALID_SOCKET)
        return;
    
    ::getsockopt(sockfd, level, option, reinterpret_cast<char*>(value), &length);
}

void Socket::setLinger(bool on, int seconds)
{
    struct linger l;
    l.l_onoff  = on ? 1 : 0;
    l.l_linger = seconds;
    setOption(SOL_SOCKET, SO_LINGER, &l, sizeof(l));
}

void Socket::getLinger(bool& on, int& seconds)
{
    struct linger l;
    unsigned int len = sizeof(l);
    getOption(SOL_SOCKET, SO_LINGER, &l, len);
    on      = l.l_onoff != 0;
    seconds = l.l_linger;
}

void Socket::setNoDelay(bool flag)
{
    int value = flag ? 1 : 0;
    setOption(IPPROTO_TCP, TCP_NODELAY, &value,sizeof(value));
}

bool Socket::getNoDelay()
{
    int value(0);
    unsigned int len = sizeof(int);
    getOption(IPPROTO_TCP, TCP_NODELAY, &value,len);
    return value != 0;
}

void Socket::setNoDelay(int sockfd,bool flag)
{
    int value = flag ? 1 : 0;
    setOption(sockfd,IPPROTO_TCP, TCP_NODELAY, &value,sizeof(value));
}

bool Socket::getNoDelay(int sockfd)
{
    int value(0);
    unsigned int len = sizeof(int);
    getOption(sockfd,IPPROTO_TCP, TCP_NODELAY, &value,len);
    return value != 0;
}

void Socket::setKeepAlive(bool flag)
{
    int value = flag ? 1 : 0;
    setOption(SOL_SOCKET, SO_KEEPALIVE, &value,sizeof(value));
}

bool Socket::getKeepAlive()
{
    int value(0);
    unsigned int len = sizeof(int);
    getOption(SOL_SOCKET, SO_KEEPALIVE, &value,len);
    return value != 0;
}

void Socket::setReuseAddress(bool flag)
{
    int value = flag ? 1 : 0;
    setOption(SOL_SOCKET, SO_REUSEADDR, &value,sizeof(value));
}

bool Socket::getReuseAddress()
{
    int value(0);
    unsigned int len = sizeof(int);
    getOption(SOL_SOCKET, SO_REUSEADDR, &value,len);
    return value != 0;
}

void Socket::setReusePort(bool flag)
{
    int value = flag ? 1 : 0;
    setOption(SOL_SOCKET, SO_REUSEPORT, &value,sizeof(value));
}

bool Socket::getReusePort()
{
    int value(0);
    unsigned int len = sizeof(int);
    getOption(SOL_SOCKET, SO_REUSEPORT, &value,len);
    return value != 0;
}

void Socket::setOOBInline(bool flag)
{
    int value = flag ? 1 : 0;
    setOption(SOL_SOCKET, SO_OOBINLINE, &value,sizeof(value));
}

bool Socket::getOOBInline()
{
    int value(0);
    unsigned int len = sizeof(int);
    getOption(SOL_SOCKET, SO_OOBINLINE, &value,len);
    return value != 0;
}

void Socket::setBroadcast(bool flag)
{
    int value = flag ? 1 : 0;
    setOption(SOL_SOCKET, SO_BROADCAST, &value,sizeof(value));
}

bool Socket::getBroadcast()
{
    int value(0);
    unsigned int len = sizeof(int);
    getOption(SOL_SOCKET, SO_BROADCAST, &value,len);
    return value != 0;
}

void Socket::setBlocking(int sockfd,bool flag)
{
    int arg = fcntl(sockfd,F_GETFL);
    long flags = arg & ~O_NONBLOCK;
    if (!flag)
        flags |= O_NONBLOCK;
    fcntl(sockfd,F_SETFL, flags);
}

bool Socket::getBlocking(int sockfd)
{
    int arg = fcntl(sockfd,F_GETFL);
    return (arg & ~O_NONBLOCK) == O_NONBLOCK;
}

int Socket::socketError()
{
    int result(0);
    unsigned int len = sizeof(int);
    getOption(SOL_SOCKET, SO_ERROR, &result,len);
    return result;
}

void Socket::ioctl(int request, int& arg)
{
    ::ioctl(m_sockfd, request, &arg);
}

void Socket::ioctl(int request, void* arg)
{
    ::ioctl(m_sockfd, request, arg);
}

int Socket::fcntl(int sockfd,int request)
{
    int rc = ::fcntl(sockfd, request);
    return rc;
}

int Socket::fcntl(int sockfd,int request, long arg)
{
    int rc = ::fcntl(sockfd, request, arg);
    return rc;
}




